from data_reader import DataReader
from image_writer import ImageWriter
from knn import KNN

class DigitClassification():
  def __init__(self,
      train_file,
      train_trans_file,
      test_file,
      test_trans_file,
      neighbour_size = 3):
    knn = KNN(test_trans_file, train_trans_file, neighbour_size)
    neighbours = knn.get_neighbours()

    img_writer = ImageWriter('working_dir/')

    train_reader = DataReader(train_file)
    test_reader = DataReader(test_file)

    train_vectors = train_reader.get_feature_vectors()
    test_vectors = test_reader.get_feature_vectors()

    test_idx = 0
    for test_vector in test_vectors:
      img_writer.create_images(test_vector, train_vectors, neighbours[test_idx])
      test_idx += 1 

digit_classifier = DigitClassification('optdigits_tra.dat',
    'optdigits_tra_trans.dat',
    'optdigits_trial.dat',
    'optdigits_trial_trans.dat')
